//===----------------------------------------------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2024 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

import Foundation
import SwiftSyntax

extension SyntaxProtocol {
  /// Returns all labeled statements available at a particular syntax node.
  ///
  /// - Returns: Available labeled statements at a particular syntax node
  /// in the exact order they appear in the source code, starting with the innermost statement.
  ///
  /// Example usage:
  /// ```swift
  /// one: while cond1 {
  ///   func foo() {
  ///     two: while cond2 {
  ///       three: while cond3 {
  ///         break // 1
  ///       }
  ///       break // 2
  ///     }
  ///   }
  ///   break // 3
  /// }
  /// ```
  /// When calling this function at the first `break`, it returns `three` and `two` in this exact order.
  /// For the second `break`, it returns only `two`.
  /// The results don't include `one`, which is unavailable at both locations due to the encapsulating function body.
  /// For `break` numbered 3, the result is `one`, as it's outside the function body and within the labeled statement.
  /// The function returns an empty array when there are no available labeled statements.
  @_spi(Experimental) public func lookupLabeledStmts() -> [LabeledStmtSyntax] {
    collectNodesOfTypeUpToFunctionBoundary(LabeledStmtSyntax.self)
  }

  /// Returns the catch node responsible for handling an error thrown at a particular syntax node.
  ///
  /// - Returns: The catch node responsible for handling an error thrown at the lookup source node.
  /// This could be a `do` statement, `try?`, `try!`, `init`, `deinit`, accessors, closures, or function declarations.
  ///
  /// Example usage:
  /// ```swift
  /// func x() {
  ///   do {
  ///     try foo()
  ///     try? bar()
  ///   } catch {
  ///     throw error
  ///   }
  /// }
  /// ```
  /// When calling this function on `foo`, it returns the `do` statement.
  /// Calling the function on `bar` results in `try?`.
  /// When used on `error`, the function returns the function declaration `x`.
  /// The function returns `nil` when there's no available catch node.
  @_spi(Experimental) public func lookupCatchNode() -> Syntax? {
    lookupCatchNodeHelper(traversedCatchClause: false)
  }

  // MARK: - lookupCatchNode

  /// Given syntax node location, finds where an error could be caught. If `traverseCatchClause` is set to `true` lookup will skip the next do statement.
  private func lookupCatchNodeHelper(traversedCatchClause: Bool) -> Syntax? {
    guard let parent else { return nil }

    switch parent.as(SyntaxEnum.self) {
    case .doStmt:
      if traversedCatchClause {
        return parent.lookupCatchNodeHelper(traversedCatchClause: false)
      } else {
        return parent
      }
    case .catchClause:
      return parent.lookupCatchNodeHelper(traversedCatchClause: true)
    case .tryExpr(let tryExpr):
      if tryExpr.questionOrExclamationMark != nil {
        return parent
      } else {
        return parent.lookupCatchNodeHelper(traversedCatchClause: traversedCatchClause)
      }
    case .functionDecl, .accessorDecl, .initializerDecl, .deinitializerDecl, .closureExpr:
      return parent
    case .exprList(let exprList):
      if let tryExpr = exprList.first?.as(TryExprSyntax.self), tryExpr.questionOrExclamationMark != nil {
        return Syntax(tryExpr)
      }
      return parent.lookupCatchNodeHelper(traversedCatchClause: traversedCatchClause)
    default:
      return parent.lookupCatchNodeHelper(traversedCatchClause: traversedCatchClause)
    }
  }

  // MARK: - walkParentTree helper methods

  /// Returns the innermost node of the specified type up to a function boundary.
  fileprivate func innermostNodeOfTypeUpToFunctionBoundary<T: SyntaxProtocol>(
    _ type: T.Type
  ) -> T? {
    collectNodesOfTypeUpToFunctionBoundary(type, stopWithFirstMatch: true).first
  }

  /// Collect syntax nodes matching the collection type up until encountering one of the specified syntax nodes. The nodes in the array are inside out, with the innermost node being the first.
  fileprivate func collectNodesOfTypeUpToFunctionBoundary<T: SyntaxProtocol>(
    _ type: T.Type,
    stopWithFirstMatch: Bool = false
  ) -> [T] {
    collectNodes(
      ofType: type,
      upTo: [
        MemberBlockSyntax.self,
        FunctionDeclSyntax.self,
        InitializerDeclSyntax.self,
        DeinitializerDeclSyntax.self,
        AccessorDeclSyntax.self,
        ClosureExprSyntax.self,
        SubscriptDeclSyntax.self,
      ],
      stopWithFirstMatch: stopWithFirstMatch
    )
  }

  /// Callect syntax nodes matching the collection type up until encountering one of the specified syntax nodes.
  private func collectNodes<T: SyntaxProtocol>(
    ofType type: T.Type,
    upTo stopAt: [SyntaxProtocol.Type],
    stopWithFirstMatch: Bool = false
  ) -> [T] {
    var matches: [T] = []
    var nextSyntax: Syntax? = Syntax(self)
    while let currentSyntax = nextSyntax {
      if stopAt.contains(where: { currentSyntax.is($0) }) {
        break
      }

      if let matchedSyntax = currentSyntax.as(T.self) {
        matches.append(matchedSyntax)
        if stopWithFirstMatch {
          break
        }
      }

      nextSyntax = currentSyntax.parent
    }

    return matches
  }
}

extension FallThroughStmtSyntax {
  /// Returns the source and destination of a `fallthrough`.
  ///
  /// - Returns: `source` as the switch case that encapsulates the `fallthrough` keyword and
  /// `destination` as the switch case that the `fallthrough` directs to.
  ///
  /// Example usage:
  /// ```swift
  /// switch value {
  /// case 2:
  ///   doSomething()
  ///   fallthrough
  /// case 1:
  ///   doSomethingElse()
  /// default:
  ///   break
  /// }
  /// ```
  /// When calling this function at the `fallthrough`, it returns `case 2` and `case 1` in this exact order.
  /// The `nil` results handle ill-formed code: there's no `source` if the `fallthrough` is outside of a case.
  /// There's no `destination` if there is no case or `default` after the source case.
  @_spi(Experimental) public func lookupFallthroughSourceAndDestintation()
    -> (source: SwitchCaseSyntax?, destination: SwitchCaseSyntax?)
  {
    guard
      let originalSwitchCase = innermostNodeOfTypeUpToFunctionBoundary(
        SwitchCaseSyntax.self
      )
    else {
      return (nil, nil)
    }

    let nextSwitchCase = lookupNextSwitchCase(at: originalSwitchCase)

    return (originalSwitchCase, nextSwitchCase)
  }

  /// Given a switch case, returns the case that follows according to the parent.
  private func lookupNextSwitchCase(at switchCaseSyntax: SwitchCaseSyntax) -> SwitchCaseSyntax? {
    guard let switchCaseListSyntax = switchCaseSyntax.parent?.as(SwitchCaseListSyntax.self) else { return nil }

    var visitedOriginalCase = false

    for child in switchCaseListSyntax.children(viewMode: .sourceAccurate) {
      if let thisCase = child.as(SwitchCaseSyntax.self) {
        if thisCase.id == switchCaseSyntax.id {
          visitedOriginalCase = true
        } else if visitedOriginalCase {
          return thisCase
        }
      }
    }

    return nil
  }
}
